{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 测试匹配结构信息"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tvm.script import relax as R\n",
    "from tvm.script import tir as T\n",
    "from tvm import relax as rx\n",
    "from tvm import relay, tir\n",
    "from tvm.relax.analysis import get_var2val\n",
    "import tvm.testing\n",
    "from tvm.relax.dpl import *"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "测试匹配时带有结构信息更新的通配符:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_wildcard_with_struct_info_updates_when_matching():\n",
    "    \"\"\"A DFPattern may be restricted to a specific StructInfo\"\"\"\n",
    "\n",
    "    pat_lhs = wildcard().has_struct_info(R.Tensor([2, 3]))\n",
    "    pat_rhs = wildcard().has_struct_info(R.Tensor([2, 3]))\n",
    "    pat = is_op(\"relax.add\")(pat_lhs, pat_rhs)\n",
    "\n",
    "    def rewriter(expr, matches):\n",
    "        lhs = matches[pat_lhs]\n",
    "        rhs = matches[pat_rhs]\n",
    "        return rx.op.multiply(lhs, rhs)\n",
    "\n",
    "    @R.function(private=True)\n",
    "    def before():\n",
    "        with R.dataflow():\n",
    "            A = R.zeros([2, 3], \"int32\")\n",
    "            B = R.ones([2, 3], \"int32\")\n",
    "            C = R.add(A, B)\n",
    "\n",
    "            R.output(C)\n",
    "        return C\n",
    "\n",
    "    @R.function(private=True)\n",
    "    def expected():\n",
    "        with R.dataflow():\n",
    "            A = R.zeros([2, 3], \"int32\")\n",
    "            B = R.ones([2, 3], \"int32\")\n",
    "            C = R.multiply(A, B)\n",
    "\n",
    "            R.output(C)\n",
    "        return C\n",
    "\n",
    "    after = rewrite_call(pat, rewriter, before)\n",
    "    tvm.ir.assert_structural_equal(expected, after)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_wildcard_with_struct_info_is_no_op_when_not_matching():\n",
    "    \"\"\"StructInfoPattern requires the StructInfo provided\n",
    "\n",
    "    Here, the pattern would match, expect that the function has\n",
    "    `R.Tensor([16,32])`, and the pattern requires `R.Tensor([2,3])`.\n",
    "    \"\"\"\n",
    "\n",
    "    pat_lhs = wildcard().has_struct_info(R.Tensor([2, 3]))\n",
    "    pat_rhs = wildcard().has_struct_info(R.Tensor([2, 3]))\n",
    "    pat = is_op(\"relax.add\")(pat_lhs, pat_rhs)\n",
    "\n",
    "    def rewriter(expr, matches):\n",
    "        lhs = matches[pat_lhs]\n",
    "        rhs = matches[pat_rhs]\n",
    "        return rx.op.multiply(lhs, rhs)\n",
    "\n",
    "    @R.function(private=True)\n",
    "    def before():\n",
    "        with R.dataflow():\n",
    "            # This R.add has the same shape as the pattern, and will\n",
    "            # be updated.\n",
    "            A = R.zeros([16, 32], \"int32\")\n",
    "            B = R.ones([16, 32], \"int32\")\n",
    "            C = R.add(A, B)\n",
    "\n",
    "            R.output(C)\n",
    "        return C\n",
    "\n",
    "    expected = before\n",
    "\n",
    "    after = rewrite_call(pat, rewriter, before)\n",
    "    tvm.ir.assert_structural_equal(expected, after)\n",
    "\n",
    "\n",
    "def test_wildcard_struct_info_for_unknown_dtype():\n",
    "    \"\"\"TensorStructInfo with unknown dtype allows any dtype\"\"\"\n",
    "\n",
    "    pat_lhs = wildcard().has_struct_info(R.Tensor([2, 3]))\n",
    "    pat_rhs = wildcard().has_struct_info(R.Tensor([2, 3]))\n",
    "    pat = is_op(\"relax.add\")(pat_lhs, pat_rhs)\n",
    "\n",
    "    def rewriter(expr, matches):\n",
    "        lhs = matches[pat_lhs]\n",
    "        rhs = matches[pat_rhs]\n",
    "        return rx.op.multiply(lhs, rhs)\n",
    "\n",
    "    @R.function(private=True)\n",
    "    def before():\n",
    "        with R.dataflow():\n",
    "            A = R.zeros([2, 3], \"int32\")\n",
    "            B = R.ones([2, 3], \"int32\")\n",
    "            C = R.add(A, B)\n",
    "\n",
    "            D = R.zeros([2, 3], \"float32\")\n",
    "            E = R.ones([2, 3], \"float32\")\n",
    "            F = R.add(D, E)\n",
    "\n",
    "            output = (C, F)\n",
    "            R.output(output)\n",
    "        return output\n",
    "\n",
    "    @R.function(private=True)\n",
    "    def expected():\n",
    "        with R.dataflow():\n",
    "            A = R.zeros([2, 3], \"int32\")\n",
    "            B = R.ones([2, 3], \"int32\")\n",
    "            C = R.multiply(A, B)\n",
    "\n",
    "            D = R.zeros([2, 3], \"float32\")\n",
    "            E = R.ones([2, 3], \"float32\")\n",
    "            F = R.multiply(D, E)\n",
    "\n",
    "            output = (C, F)\n",
    "            R.output(output)\n",
    "        return output\n",
    "\n",
    "    after = rewrite_call(pat, rewriter, before)\n",
    "    tvm.ir.assert_structural_equal(expected, after)\n",
    "\n",
    "\n",
    "def test_wildcard_struct_info_with_symbolic_vars():\n",
    "    \"\"\"StructInfoPattern may define symbolic vars\n",
    "\n",
    "    This test finds an elementwise `R.add`, while ignoring a\n",
    "    broadcasted `R.add`.\n",
    "    \"\"\"\n",
    "\n",
    "    m = tir.Var(\"m\", \"int64\")\n",
    "    n = tir.Var(\"n\", \"int64\")\n",
    "\n",
    "    pat_lhs = wildcard().has_struct_info(R.Tensor([m, n]))\n",
    "    pat_rhs = wildcard().has_struct_info(R.Tensor([m, n]))\n",
    "    pat = is_op(\"relax.add\")(pat_lhs, pat_rhs)\n",
    "\n",
    "    def rewriter(expr, matches):\n",
    "        lhs = matches[pat_lhs]\n",
    "        rhs = matches[pat_rhs]\n",
    "        return rx.op.multiply(lhs, rhs)\n",
    "\n",
    "    @R.function(private=True)\n",
    "    def before():\n",
    "        with R.dataflow():\n",
    "            A = R.zeros([64, 128], \"int32\")\n",
    "            B = R.ones([64, 128], \"int32\")\n",
    "            C = R.add(A, B)\n",
    "\n",
    "            D = R.zeros([64, 128], \"float32\")\n",
    "            E = R.ones([1, 128], \"float32\")\n",
    "            F = R.add(D, E)\n",
    "\n",
    "            output = (C, F)\n",
    "            R.output(output)\n",
    "        return output\n",
    "\n",
    "    @R.function(private=True)\n",
    "    def expected():\n",
    "        with R.dataflow():\n",
    "            A = R.zeros([64, 128], \"int32\")\n",
    "            B = R.ones([64, 128], \"int32\")\n",
    "            C = R.multiply(A, B)\n",
    "\n",
    "            D = R.zeros([64, 128], \"float32\")\n",
    "            E = R.ones([1, 128], \"float32\")\n",
    "            F = R.add(D, E)\n",
    "\n",
    "            output = (C, F)\n",
    "            R.output(output)\n",
    "        return output\n",
    "\n",
    "    after = rewrite_call(pat, rewriter, before)\n",
    "    tvm.ir.assert_structural_equal(expected, after)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ai",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
